import pandas as pd
import numpy as np
import torch
import json
import cv2
import os
import torch
import time
from tqdm import tqdm
import torch.nn as nn
# from progressbar import ProgressBar
from data_aug import resize,center_crop,random_crop,aug_predict,multi_model
from models.model import model_resnet_unet_GN,model_hrnet,model_Unet_efficient_b3
from make_heatmap import heatmap_to_anno,group_heatmap_to_anno
from dataset import resize_point

def compute_loss(pred,label,original_size = None,resize_size = None,distance = 1,spacing = 0.143):
    # 传入pred和lable的单张图片的坐标
    ## 用来记录距离小于指定数值的正样本个数
    positive = 0
    ## 用来记录l2 loss
    l2_loss = []
    assert len(pred) == len(label),'the len of pred and label shoule be same'
    ## 输入全部的pred和label，用list中包含tuple的形式，其中像素间距为0.139mm
    for index in range(len(pred)):
        single_pred = pred[index]
        single_label = label[index]
        ## 遇到缺失的label不进行计算
        if 0 in single_label:
            continue
        ## 计算两个点之间的像素距离。注意根据原始尺寸和resize后的尺寸对坐标值进行转换。
        if original_size is not None:
            # loss = (((single_pred[0]-single_label[0])/resize_size[0]*original_size[0])**2 +
            #         ((single_pred[1]-single_label[1])/resize_size[1]*original_size[1])** 2)**0.5
            loss = abs(single_pred[0] - single_label[0]) / resize_size[0] * original_size[0]
        else:
            ## 如果使用原图上的坐标。直接进行计算即可，不需要进行转换坐标。
            # loss = (((single_pred[0] - single_label[0])) ** 2 +
            #         ((single_pred[1] - single_label[1])) ** 2) ** 0.5
            loss = abs(single_pred[0] - single_label[0])
        ## 将计算出的像素距离转换成物理距离
        distance_loss = float(loss)*spacing
        # print('distance loss',distance_loss)
        l2_loss.append(distance_loss)
        if distance_loss <= distance:
            positive +=1
    return positive,l2_loss





def val_fun(model,data_path,json_dir,image_size,distance = 1.,spacing = 0.143,show_pred = False,
            test_aug = False,print_l2_loss = False,):
    json_data = json.load(open(json_dir))
    save_json = []
    # print(name_list)
    all_l2_loss = []
    all_positive = 0
    with torch.no_grad():
        model.eval()
        # progressbar = ProgressBar()
        for p in tqdm(range(len(json_data))):
            image_dir = os.path.join(data_path,json_data[p]['name'] + '.jpg')
            image = cv2.imread(image_dir)
            original_size = image.shape[:-1]
            ## 读取label
            label = json_data[p]['point']
            label = [p[0][::-1] for p in label]
            if image_size is not None:
                try:
                    image, label = resize_point(image, label, image_size)
                except Exception as e:
                    print(e)
                    continue
                
            # image = np.flip(image,1).copy()

            input = torch.from_numpy(np.transpose(image,axes=(2,0,1))).float().cuda()
            input = input[np.newaxis,...]
            
            if test_aug:
                pred = aug_predict(model,input)
            else:
                pred = model(input)

            
            ## 注意pred[0]去掉batch维度
            pred = heatmap_to_anno(pred[0].cpu())
            
            # cv2.imshow('heatmap of predict',np.sum(pred[0].cpu().numpy(),axis=0).astype(np.uint8))
            # cv2.waitKey()

            ## 分组预测
            # pred = group_heatmap_to_anno(pred[0].cpu())
            # # 胸椎棘突
            # pred1 = pred[:4]
            # # 肋骨
            # pred2 = pred[4:12]
            # # 肩部
            # pred3 = pred[12:20]
            # # 肺部
            # pred4 = pred[20:]
            
            
            # pred = [pred[0]]
            # label = [label[0]]
            pred_ori = [[p[0]/image_size[0] * original_size[0],p[1] / image_size[1] * original_size[1]] for p in pred]
            save_json.append({'name':json_data[p]['name'],'point':[[pred_ori[index][::-1],index] for index in range(len(pred))]})
            with open('./label_file/calculate_result/pred_val.json', 'w') as f:
                json.dump(save_json, f)
            # print(save_json)
            positive,l2_loss = compute_loss(pred,label,original_size=None,resize_size=None,distance=distance,spacing=spacing)
            # print(l2_loss)
            if show_pred:
                pass
                print('--l2 loss--',l2_loss,'----',np.mean(l2_loss))
            all_positive += positive
            all_l2_loss = all_l2_loss+l2_loss
            
            if show_pred:
                # for index in range(len(pred)):
                #     single_pred = pred[index]
                #     single_label = label[index]
                #     single_label = (int(single_label[0]),int(single_label[1]))
                #     if 0 in single_label :
                #         continue
                #     ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                #     cv2.circle(image, single_label[::-1],2,(255,128,255),thickness=6)
                #     cv2.circle(image, single_pred[::-1], 1, (128, 255, 0), thickness=6)
                #     # cv2.putText(image, name_list[index], tuple(map(lambda x:x+15,single_pred[::-1])), cv2.FONT_HERSHEY_PLAIN, 1, (128, 255, 0), 1)
                
                
                
                for index in range(len(label)):
                    single_label = label[index]
                    single_label = (int(single_label[0]),int(single_label[1]))
                    if 0 in single_label :
                        continue
                    ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                    cv2.circle(image, single_label[::-1],1,(255,128,255),thickness=3)

                for index in range(len(pred)):
                    single_pred = pred[index]
                    if 0 in single_label :
                        continue
                    ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                    cv2.circle(image, single_pred[::-1], 1, (128, 255, 0), thickness=3)
                
                
                
                ## 分组训练预测，不同组的点用不同颜色显示。
                # for index in range(len(pred1)):
                #     single_pred = pred1[index]
                #     if 0 in single_label :
                #         continue
                #     ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                #     cv2.circle(image, single_pred[::-1], 1, (128, 255, 0), thickness=6)

                # for index in range(len(pred2)):
                #     single_pred = pred2[index]
                #     if 0 in single_label :
                #         continue
                #     ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                #     cv2.circle(image, single_pred[::-1], 1, (200, 0, 128), thickness=6)
                #
                #
                #
                # for index in range(len(pred3)):
                #     single_pred = pred3[index]
                #     if 0 in single_label :
                #         continue
                #     ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                #     cv2.circle(image, single_pred[::-1], 1, (128, 0, 255), thickness=6)
                #
                #
                # for index in range(len(pred4)):
                #     single_pred = pred4[index]
                #     if 0 in single_label :
                #         continue
                #     ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                #     cv2.circle(image, single_pred[::-1], 1, (0, 128, 255), thickness=6)
                    
                    
                
                
                if sum(i > 5 for i in l2_loss) > 0:
                    # print(image.shape)
                    # cv2.imwrite('./haha/'+json_data[p]['name'] + '.jpg',image)
                    cv2.imshow('image',image)
                    cv2.waitKey(0)
                # cv2.imwrite(os.path.join('../data/process_data/show/pred',image_dir.split('/')[-1]), image)
            if print_l2_loss:
                show_bad = 0
                for i in l2_loss:
                    if i>5:
                        show_bad += 1
                if show_bad > 0:
                    print(l2_loss)
                    for index in range(len(pred)):
                        single_pred = pred[index]
                        single_label = label[index]
                        single_label = (int(single_label[0]), int(single_label[1]))
                        if 0 in single_label:
                            continue
                        ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                        cv2.circle(image, single_label[::-1], 1, (255, 128, 255), thickness=8)
                        cv2.circle(image, single_pred[::-1], 1, (128, 255, 0), thickness=8)

    
                    cv2.imshow('image', image)
                    cv2.waitKey(0)
    # csv_data.to_csv('../data/pred_result/baseline.csv',index=False)
    # 由于部分label存在缺失，不进行loss的计算。根据all_l2_loss的长度来作为实际参与计算的点的数量。
    
    return all_positive/len(all_l2_loss),np.mean(all_l2_loss),np.median(all_l2_loss)


def get_loss_form_json(pred_json='./label_file/calculate_result/pred_val.json',
                       label_json='label_file_0304/point_0304_val.json',
                       spacing=0.545,distance=1):
    all_l2_loss = []
    all_positive = 0
    pred_list = json.load(open(pred_json))
    label_list = json.load(open(label_json))
    for index in range(len(label_list)):
        label = label_list[index]['point']
        label = [p[0][::-1] for p in label]
        pred = pred_list[index]['point']
        pred = [p[0][::-1] for p in pred]
        
        pred = [pred[3]]
        label = [label[3]]
        positive, l2_loss = compute_loss(pred, label, original_size=None, resize_size=None, distance=distance,
                                         spacing=spacing)
        all_positive += positive
        all_l2_loss = all_l2_loss + l2_loss
    print(all_positive / len(all_l2_loss), np.mean(all_l2_loss), np.median(all_l2_loss))


if __name__ == '__main__':
    get_loss_form_json(distance=1)
    
    # # for i in range(31,249,2):
    # #     i = '199'
    #
    #     start_time = time.time()
    #
    #     # Image_Size = (640,512)
    #     Image_Size = (768,576)
    #
    #
    #     # val_csv_dir = '../data/csv_file/val_data.csv'
    #     Data_Path = './data_0304/images'
    #     Json_Dir = 'label_file_0304/point_0304_val.json'
    #     # model = model_cpn50(image_size = Image_Size)-
    #     # model = model_simple_pose_res101()
    #     model = model_resnet_unet_GN(layer='resnet34', use_aspp=True, num_classes=10, pre_train=False)
    #     model = nn.DataParallel(model)
    #     # model.load_state_dict(torch.load('../model_save/baseline/149.pth'))
    #     # model.load_state_dict(torch.load('../model_save/Group_1230/291.pth'))
    #     # model.load_state_dict(torch.load('../model_save/Group_0108/251.pth'))
    #     model.load_state_dict(torch.load('../model_save/0304/baseline/141.pth'))
    #     # model.load_state_dict(torch.load('../model_save/Group_1230/11.pth'))
    #     recall,l2_mean,l2_medium = val_fun(model,data_path=Data_Path,json_dir=Json_Dir,image_size=Image_Size,
    #             distance=10,spacing=0.545,show_pred = False,test_aug=False,print_l2_loss=False,)
    #     print('--recall',recall,'l2 mean',l2_mean,'l2 medium',l2_medium)
    #     # print('use time ---',time.time() - start_time)